import numpy as np
import matplotlib.pyplot as plt
import os


def model(X, theta):
    return X.dot(theta)


def sigmoid(z):
    return 1 / (1 + np.exp(-z))


def cost_func(h, y):
    m = len(h)
    return - 1 / m * np.sum(y * np.log(h) + (1 - y) * np.log(1 - h))


def FP(X, theta1, theta2, theta3):
    a1 = X
    z2 = model(a1, theta1)
    a2 = sigmoid(z2)
    z3 = model(a2, theta2)
    a3 = sigmoid(z3)
    z4 = model(a3, theta3)
    a4 = sigmoid(z4)
    return a2, a3, a4


def BP(X, y, theta1, theta2, theta3, a2, a3, a4, alpha):
    m  = len(X)
    s4 = a4 - y
    s3 = s4.dot(theta3.T) * (a3 * (1 - a3))
    s2 = s3.dot(theta2.T) * (a2 * (1 - a2))
    dt3 = 1 / m * a3.T.dot(s4)   # ATTENTION: 1 / m
    dt2 = 1 / m * a2.T.dot(s3)
    dt1 = 1 / m * X.T.dot(s2)
    theta3 -= alpha * dt3
    theta2 -= alpha * dt2
    theta1 -= alpha * dt1
    return theta1, theta2, theta3


def grad(X, y, alpha=0.01, iter0=2000):
    L2, L3, L4 = 100, 64, 10
    m, n = X.shape
    group = iter0 // 20
    theta1 = np.random.randn(n, L2)
    theta2 = np.random.randn(L2, L3)
    theta3 = np.random.randn(L3, L4)
    j_his = np.zeros(iter0)
    for i in range(iter0):
        a2, a3, a4 = FP(X, theta1, theta2, theta3)
        j = cost_func(a4, y)
        j_his[i] = j
        if 0 == i % group:
            print(f'#{i + 1} cost func value = {j}')
        theta1, theta2, theta3 = BP(X, y, theta1, theta2, theta3, a2, a3, a4, alpha)
    if 0 != i % group:
        print(f'#{i + 1} cost func value = {j}')
    return theta1, theta2, theta3, j_his, a4


def score(h, y):
    h_ = h.argmax(axis=1)
    y_ = y.argmax(axis=1)
    return np.mean(y_ == h_)


if '__main__' == __name__:
    np.random.seed(1)
    plt.figure(figsize=[16, 12])
    spr = 3  # subplot row
    spc = 6  # subplot column
    spn = 0

    # load data
    path = r'../../../../large_data/ML1/hand_writing/'
    data = np.loadtxt(path + 'imgX.txt', delimiter=',')
    y = np.loadtxt(path + 'labely.txt', delimiter=',')
    m = len(y)
    y[y == 10] = 0  # ATTENTION

    # scale data
    mu = data.mean()
    sigma = data.std()
    data -= mu
    data /= sigma

    # shuffle data
    rnd_idx = np.random.permutation(m)
    data = data[rnd_idx]
    y = y[rnd_idx]

    # splice data
    X = np.c_[np.ones(m), data]

    # onehot
    y_onehot = np.zeros([m, 10])
    for k, v in enumerate(y):
        y_onehot[k, int(v)] = 1
    # print(y[:10])
    # print(y_onehot[:10])

    # split data
    m_train = int(0.7 * m)
    m_test = m - m_train
    XX_train, XX_test = np.split(X, [m_train])
    x_train, x_test = np.split(data, [m_train])
    y_train, y_test = np.split(y, [m_train])
    y_onehot_train, y_onehot_test = np.split(y_onehot, [m_train])

    # train
    ver = '1.0'
    alpha = 0.1
    iter0 = 2000
    save_path = os.path.basename(__file__) + '_' + ver + '_' + str(alpha) + '_' + str(iter0) + '_theta'
    suffix = '.tmp.txt'
    if os.path.exists(save_path + '1' + suffix):
        print('Cache file exists, will read thetas in cache file...')
        theta1 = np.loadtxt(save_path + '1' + suffix)
        theta2 = np.loadtxt(save_path + '2' + suffix)
        theta3 = np.loadtxt(save_path + '3' + suffix)
    else:
        theta1, theta2, theta3, j_his, h_train = grad(XX_train, y_onehot_train, alpha, iter0)
        np.savetxt(save_path + '1' + suffix, theta1)
        np.savetxt(save_path + '2' + suffix, theta2)
        np.savetxt(save_path + '3' + suffix, theta3)
        print(f'Training score = {score(h_train, y_onehot_train)}')
        spn += 1
        plt.subplot(spr, spc, spn)
        plt.plot(j_his, label='cost function value')
        plt.grid()
        plt.legend()

    # test
    a2_test, a3_test, h_test = FP(XX_test, theta1, theta2, theta3)
    print(f'Testing score = {score(h_test, y_onehot_test)}')


    # show detail
    def output(tgt, num_all):
        global spn
        num = 0
        idx = -1
        while True:
            idx += 1
            if y_test[idx] != tgt:
                continue
            num += 1
            if num > num_all:
                break
            spn += 1
            plt.subplot(spr, spc, spn)
            plt.title(str(y_test[idx]) + ' => ' + str(h_test[idx].argmax()))
            plt.imshow(x_test[idx].reshape(20, 20).T)
            spn += 1
            plt.subplot(spr, spc, spn)
            plt.imshow(a2_test[idx].reshape(10, 10).T)
            spn += 1
            if spn > spr * spc:
                break
            plt.subplot(spr, spc, spn)
            plt.imshow(a3_test[idx].reshape(8, 8).T)


    output(0, 3)
    output(1, 3)

    plt.show()
